package parser_antlr4.doris;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import constant.SQLParserConstants;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import domain.Column;
import domain.Table;
import domain.property.DorisProperty;
import domain.property.DynamicPartition;
import org.antlr.v4.runtime.CommonTokenStream;
import org.antlr.v4.runtime.Token;
import org.antlr.v4.runtime.misc.Interval;
import org.antlr.v4.runtime.tree.TerminalNode;
import parser_antlr4.doris.antlr4.DorisParser;
import parser_antlr4.doris.antlr4.DorisParserBaseListener;
import utils.StringSQLUtils;

import java.util.*;

@Slf4j
public class DorisSQLListener extends DorisParserBaseListener {
  private static final String EXTERNAL = "EXTERNAL";
  private static final String DB_TYPE = "doris";

  private CommonTokenStream commonTokenStream;

  public DorisSQLListener(CommonTokenStream commonTokenStream) {
    this.commonTokenStream = commonTokenStream;
  }

  @Getter
  private final List<Table> tableList = new ArrayList<>(10);

  private Table table;
  private List<Column> columnList;
  private DorisProperty dorisProperty;
  private DynamicPartition dynamicPartition;

  @Override
  public void enterCreateTableStatement(DorisParser.CreateTableStatementContext ctx) {
    table = new Table();
    columnList = new ArrayList<>();
    dorisProperty = new DorisProperty();
    dynamicPartition = new DynamicPartition();

    table.setDbType(DB_TYPE);
    Interval sourceInterval = ctx.getSourceInterval();
    table.setOriginDDL(getSQLSegment(sourceInterval.a, sourceInterval.b));
    parseTableNameAndSchema(ctx);
    parseTableComment(ctx);
    parseExternal(ctx);
    parseEngine(ctx);
    parseKeysType(ctx);
    parseStaticPartition(ctx);
    parseBuckets(ctx);
    parseRollup(ctx);
    parseDynamicPartitionAndProperties(ctx);
    parseExtProperties(ctx);
  }

  @Override
  public void exitCreateTableStatement(DorisParser.CreateTableStatementContext ctx) {
    table.setColumnList(columnList);
    dorisProperty.setDynamicPartition(dynamicPartition);
    table.setProperties(JSONUtil.toJsonStr(dorisProperty));
    tableList.add(table);
  }

  void parseTableNameAndSchema(DorisParser.CreateTableStatementContext ctx) {
    if (ctx.name != null) {
      String fullTableName = StringSQLUtils.cleanBackQuote(ctx.name.getText());
      String[] schemaAndName = fullTableName.split("\\.");
      if (schemaAndName.length == 1) {
        table.setTableName(schemaAndName[0]);
      } else {
        table.setTableName(schemaAndName[0]);
        table.setSchema(schemaAndName[1]);
      }
    }
  }

  void parseTableComment(DorisParser.CreateTableStatementContext ctx) {
    TerminalNode comment = ctx.COMMENT();
    if (comment != null) {
      TerminalNode terminalNode = ctx.STRING_LITERAL();
      table.setTableComment(StringSQLUtils.cleanDoubleQuote(terminalNode.getText()));
    }
  }

  void parseExternal(DorisParser.CreateTableStatementContext ctx) {
    TerminalNode external = ctx.EXTERNAL();
    if (external != null) {
      dorisProperty.setExternal(EXTERNAL.equalsIgnoreCase(external.getText()));
    }
  }

  void parseEngine(DorisParser.CreateTableStatementContext ctx) {
    TerminalNode engine = ctx.ENGINE();
    if (engine != null) {
      dorisProperty.setEngine(ctx.engine.getText());
    }
  }

  void parseKeysType(DorisParser.CreateTableStatementContext ctx) {
    Integer keyTypeStart = null;
    Integer keyTypeEnd = null;
    String keysType = null;

    if (ctx.AGGREGATE() != null) {
      keyTypeStart = ctx.AGGREGATE().getSourceInterval().a;
      keysType = ctx.AGGREGATE().getText();
    }

    if (ctx.UNIQUE() != null) {
      keyTypeStart = ctx.UNIQUE().getSourceInterval().a;
      keysType = ctx.UNIQUE().getText();
    }

    if (ctx.DUPLICATE() != null) {
      keyTypeStart = ctx.DUPLICATE().getSourceInterval().a;
      keysType = ctx.DUPLICATE().getText();
    }

    dorisProperty.setKeysType(keysType);
    if (keyTypeStart != null) {
      keyTypeEnd = ctx.keys.getSourceInterval().b;
      if (ctx.CLUSTER() != null) {
        keyTypeEnd = ctx.clusterKeys.getSourceInterval().b;
      }
      dorisProperty.setOriginKeySQL(getSQLSegment(keyTypeStart, keyTypeEnd));
    }
  }

  void parseStaticPartition(DorisParser.CreateTableStatementContext ctx) {
    DorisParser.PartitionTableContext partition = ctx.partition;
    dorisProperty.setPartitionType(DorisProperty.PARTITION_STATIC);
    if (partition != null) {
      Interval sourceInterval = partition.getSourceInterval();
      dorisProperty.setPartition(getSQLSegment(sourceInterval.a, sourceInterval.b));
    }
  }

  void parseBuckets(DorisParser.CreateTableStatementContext ctx) {
    TerminalNode distributed = ctx.DISTRIBUTED();
    if (distributed != null) {
      int bucketsStart = distributed.getSourceInterval().a;
      Integer bucketsEnd = null;

      if (ctx.HASH() != null) {
        bucketsEnd = ctx.HASH().getSourceInterval().b;
      }
      if (ctx.RANDOM() != null) {
        bucketsEnd = ctx.RANDOM().getSourceInterval().b;
      }
      if (ctx.BUCKETS() != null) {
        if (ctx.INTEGER_VALUE() != null) {
          bucketsEnd = ctx.INTEGER_VALUE().getSourceInterval().b;
        }
        if (ctx.AUTO() != null) {
          bucketsEnd = ctx.AUTO().getSourceInterval().b;
        }
      }

      if (bucketsEnd != null) {
        dorisProperty.setBucket(getSQLSegment(bucketsStart, bucketsEnd));
      }
    }
  }

  void parseRollup(DorisParser.CreateTableStatementContext ctx) {
    TerminalNode rollup = ctx.ROLLUP();
    if (rollup != null) {
      int rollupStart = rollup.getSourceInterval().a;
      int rollupEnd = ctx.rollupDefs().getSourceInterval().b;
      String rollupSQL = getSQLSegment(rollupStart, rollupEnd);
      rollupSQL = rollupSQL + ")";
      dorisProperty.setRollupSQL(rollupSQL);
    }
  }

  void parseDynamicPartitionAndProperties(DorisParser.CreateTableStatementContext ctx) {
    DorisParser.PropertyClauseContext propCtx = ctx.properties;
    Map<String, String> properties = new LinkedHashMap<>();

    if (propCtx != null) {
      for (DorisParser.PropertyItemContext property : propCtx.propertyItemList().properties) {
        String key = StringSQLUtils.cleanDoubleQuote(property.key.getText());
        String value = StringSQLUtils.cleanDoubleQuote(property.value.getText());
        properties.put(key, value);
      }

      boolean isDynamicPartition = Boolean.parseBoolean(properties.get(DynamicPartition.K_ENABLE));
      if (isDynamicPartition) {
        dorisProperty.setPartition(null);
        dorisProperty.setPartitionType(DorisProperty.PARTITION_DYNAMIC);
        DorisParser.PartitionTableContext partition = ctx.partition;
        if (partition != null) {
          String text = partition.identityOrFunctionList().getText();
          if (StrUtil.isNotEmpty(text)) {
            if (text.startsWith("(")) {
              text = text.substring(1);
            }
            if (text.endsWith(")")) {
              text = text.substring(0, text.length() - 1);
            }
            dynamicPartition.setColumn(text);
          }
        }
        dynamicPartition.setTimeUnit(DynamicPartition.TimeUnit.valueOf(properties.get(DynamicPartition.K_TIME_UNIT)));
        dynamicPartition.setPrefix(properties.get(DynamicPartition.K_PREFIX));
        dynamicPartition.setEnd(Integer.parseInt(properties.get(DynamicPartition.K_END)));
      }

      String replicationNumStr = properties.get(DorisProperty.K_REPLICATION_NUM);
      if (Objects.isNull(replicationNumStr)) {
        replicationNumStr = properties.get(DorisProperty.K_REPLICATION_ALLOCATION);
        if (Objects.nonNull(replicationNumStr) && replicationNumStr.contains(":")) {
          replicationNumStr = replicationNumStr.split(":")[1];
        }
      }

      if (replicationNumStr == null) {
        throw new RuntimeException(SQLParserConstants.REPLICATION_NUM_BLANK);
      }

      try {
        int replicationNum = Integer.parseInt(replicationNumStr.trim());
        dorisProperty.setReplicationNum(replicationNum);
      } catch (NumberFormatException e) {
        throw new RuntimeException(String.format(SQLParserConstants.REPLICATION_NUM_NUMBER, replicationNumStr));
      }

      dorisProperty.setDynamicPartition(dynamicPartition);
      dorisProperty.setProperties(properties);
    }
  }

  void parseExtProperties(DorisParser.CreateTableStatementContext ctx) {
    DorisParser.PropertyClauseContext propCtx = ctx.extProperties;
    Map<String, String> extProperties = new LinkedHashMap<>();

    if (propCtx != null) {
      for (DorisParser.PropertyItemContext property : propCtx.propertyItemList().properties) {
        String key = StringSQLUtils.cleanDoubleQuote(property.key.getText());
        String value = StringSQLUtils.cleanDoubleQuote(property.value.getText());
        extProperties.put(key, value);
      }
    }
    dorisProperty.setExtProperties(extProperties);
  }

  public String getSQLSegment(Integer start, Integer end) {
    StringBuilder SQL = new StringBuilder();
    List<Token> tokens = commonTokenStream.getTokens(start, end);
    for (Token token : tokens) {
      SQL.append(token.getText());
    }
    return SQL.toString();
  }
}
